package com.example.ai.output;

import com.example.ai.embedding.SimilarSearchReq;
import com.example.ai.output.converter.MyListOutputConverter;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.ai.converter.ListOutputConverter;
import org.springframework.ai.converter.MapOutputConverter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.convert.support.DefaultConversionService;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
@RestController
@RequestMapping("/structed/output")
public class StructedOutputController {
    @Autowired
    VectorStore vectorStore;
    @Autowired
    ChatModel chatModel;
    @Autowired
    ObjectMapper objectMapper;

    @PostMapping("/json")
    public ResponseEntity<?> jsonOutput(@RequestBody SimilarSearchReq req) {
        SearchRequest searchRequest = SearchRequest.builder()
                .query(req.getContent())
                .topK(req.getTopK())
                .similarityThreshold(req.getThreshold())
                .filterExpression(req.getFilter())
                .build();
        QuestionAnswerAdvisor questionAnswerAdvisor = QuestionAnswerAdvisor.builder(vectorStore)
                .searchRequest(searchRequest)
                .build();

        org.springframework.ai.chat.client.ResponseEntity responseEntity = ChatClient.builder(chatModel)
                .build().prompt()
                .advisors(new SimpleLoggerAdvisor())
                .advisors(questionAnswerAdvisor)
                .user(req.getContent())
                .call()
                        .responseEntity(new BeanOutputConverter<>(HashMap.class));
        log.info("responseEntity:{}", responseEntity);
        return ResponseEntity.ok(responseEntity);
    }

    @PostMapping("/structed/output1")
    public ResponseEntity<?> structedOutput1(@RequestBody NLReq req) {
        ActorsFilms actorsFilms = ChatClient.create(chatModel).prompt()
                .user(u -> u.text("""
                                请给与3部演员：{actor} 的电影作品
                                要求：
                                - 没有找到演员：{actor}，就直接输出空格。
                                - 没有找到演员：{actor}的电影作品，也输出空格。
                                - 不要推测和捏造电影作品
                                """)
                        .param("actor", req.getContent()))
                .advisors(new SimpleLoggerAdvisor())
                .call()
                .entity(ActorsFilms.class);
        return ResponseEntity.ok(actorsFilms);
    }

    @PostMapping("/structed/output2")
    public ResponseEntity<?> structedOutput2(@RequestBody NLReq req) {
        MapOutputConverter mapOutputConverter = new MapOutputConverter();
        Map<String,Object> map = ChatClient.create(chatModel).prompt()
                .user(u -> u.text("""
                                请给与3部演员：{actor} 的电影作品
                                要求：
                                - 没有找到演员：{actor}，就直接输出空格。
                                - 没有找到演员：{actor}的电影作品，也输出空格。
                                - 不要推测和捏造电影作品
                                """)
                        .param("actor", req.getContent()))
                .advisors(new SimpleLoggerAdvisor())
                .call()
                .entity(mapOutputConverter);
        return ResponseEntity.ok(map);
    }

    @PostMapping("/structed/output3")
    public ResponseEntity<?> structedOutput3(@RequestBody NLReq req) {
        ListOutputConverter listOutputConverter = new ListOutputConverter(new DefaultConversionService());
        List<String> listString = ChatClient.create(chatModel).prompt()
                .user(u -> u.text("""
                                请给与3部演员：{actor} 的电影作品
                                要求：
                                - 没有找到演员：{actor}，就直接输出空格。
                                - 没有找到演员：{actor}的电影作品，也输出空格。
                                - 不要推测和捏造电影作品
                                """)
                        .param("actor", req.getContent()))
                .advisors(new SimpleLoggerAdvisor())
                .call()
                .entity(listOutputConverter);
        return ResponseEntity.ok(listString);
    }


    @PostMapping("/structed/listOutput")
    public ResponseEntity<?> listOutput(@RequestBody NLReq req) throws JsonProcessingException {
        return ResponseEntity.ok(listOutputConverter(chatModel,req.getContent()));
    }

    @PostMapping("/structed/mapOutput")
    public ResponseEntity<?> mapOutput(@RequestBody NLReq req) throws JsonProcessingException {
        return ResponseEntity.ok(mapOutputConverter(chatModel,req.getContent()));
    }

    @PostMapping("/structed/listObjectOutput")
    public ResponseEntity<?> listObjectOutput(@RequestBody NLReq req) throws JsonProcessingException {
        MyListOutputConverter<ActorsFilms> listOutputConverter = new MyListOutputConverter(ActorsFilms.class);
        List<ActorsFilms> films = ChatClient.create(chatModel).prompt()
                .user(u -> u.text("""
                                请给与3部演员：{actor} 的电影作品
                                要求：
                                - 没有找到演员：{actor}，就直接输出空格。
                                - 没有找到演员：{actor}的电影作品，也输出空格。
                                - 不要推测和捏造电影作品
                                """)
                        .param("actor", req.getContent()))
                .advisors(new SimpleLoggerAdvisor())
                .call()
                .entity(listOutputConverter);
        return ResponseEntity.ok(listOutputConverter(chatModel,req.getContent()));
    }

    public List<String> listOutputConverter(ChatModel chatModel,String subject) throws JsonProcessingException {
        ListOutputConverter listConverter = new ListOutputConverter(new DefaultConversionService());

        String format = listConverter.getFormat();
        String template = """
        Provide a list of 3 {subject}
        {format}
        """;
        PromptTemplate promptTemplate = PromptTemplate.builder()
                .template(template)
                .variables(Map.of("subject", subject, "format", format))
                .build();
        Prompt prompt = new Prompt(promptTemplate.createMessage());
        log.info("Prompt : {}",objectMapper.writeValueAsString(prompt));
        Generation generation = chatModel.call(prompt).getResult();
        log.info(generation.getOutput().getText());
        List<String> items = listConverter.convert(generation.getOutput().getText());
        return items;
    }


    public Map<String, Object> mapOutputConverter(ChatModel chatModel,String subject) throws JsonProcessingException {
        MapOutputConverter mapConverter = new MapOutputConverter();

        String format = mapConverter.getFormat();
        String template = """
        Provide a dictionary of {subject}
        {format}
        """;
//        subject = "animal types and their average lifespan";
        PromptTemplate promptTemplate = PromptTemplate.builder()
                .template(template)
                .variables(Map.of("subject", subject, "format", format))
                .build();
        Prompt prompt = new Prompt(promptTemplate.createMessage());
        log.info("Prompt : {}",objectMapper.writeValueAsString(prompt));
        Generation generation = chatModel.call(prompt).getResult();
        log.info(generation.getOutput().getText());
        Map<String, Object> map = mapConverter.convert(generation.getOutput().getText());
        return map;
    }
    record ActorsFilms(String actor, List<Movie> movies) {

        record Movie(String title, String year, String daoyan,String desc) {
        }
    }
 }
